// Copyright 2023 Huawei Cloud Computing Technology Co., Ltd.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "MsgReassemble.h"

#include <cstring>

#include "Include/msghead.h"
#include "Log/logging.h"
#include "MemoryManagement/VmiBuffer.h"
#include "HWSecure/include/securec.h"

/**
 * @brief: constructor
 * @param [in] type: type
 */
MsgReassemble::MsgReassemble(VMIMsgType type)
{
    m_type = type;
    m_totalSize = 0;
}

/**
 * @brief: deconstructor, clear queue and free memory
 */
MsgReassemble::~MsgReassemble()
{
    Clear();
}

/**
 * @brief: clear queue and free memory
 */
void MsgReassemble::Clear()
{
    void* packet = nullptr;

    while (!(m_reassembleQueue.empty())) {
        packet = m_reassembleQueue.front();
        m_reassembleQueue.pop_front();
        VmiFreeBuffer(packet);
    }

    m_totalSize = 0;
}

/*
 * @brief: Set whether to send the whole package without going through the fragmentation process.
 * @param [in] flag：True means that the whole packet is sent regardless of the data.
 */
void MsgReassemble::SetSingleSend(bool flag)
{
    INFO("set singleSend flag %d", flag);
    m_singleSend = flag;
}

/**
 * @brief: check packet
 * @param [in] packet: packet
 * @return: true if ok
 */
bool MsgReassemble::CheckPacket(StreamMsgHead& packet)
{
    if (packet.getType() != m_type) {
        ERR("REASSEMBLE msgtype(%u) is error, this type is %u", packet.getType(), m_type);
        VmiFreeBuffer(&packet);
        return false;
    }

    if (packet.getPayloadSize() > FRAGMENT_SIZE && (!m_singleSend || packet.getFlag() != SINGLE_FRAGMENT)) {
        ERR("REASSEMBLE fragment error, payloadSize:(%u), FRAGMENT_SIZE:(%u),  stream_type:%u, flag:%u",
            packet.getPayloadSize(), FRAGMENT_SIZE, m_type, packet.getFlag());
        VmiFreeBuffer(&packet);
        return false;
    }
    return true;
}

/**
 * @brief: handle single fragment
 * @param [in] packet: packet
 * @return: this packet
 */
std::pair<uint8_t*, uint32_t> MsgReassemble::ProcessSingleFragment(StreamMsgHead& packet)
{
    std::pair<uint8_t*, uint32_t> nullPair(nullptr, 0);
    if (!(m_reassembleQueue.empty())) {
        ERR("REASSEMBLE error, drop some fragment(size %u), this type is %u", m_totalSize, m_type);
        Clear();
    }
    uint32_t payloadSize = packet.getPayloadSize();
    StreamMsgHead* head = &packet;

    uint8_t* message = reinterpret_cast<uint8_t *>(VmiAllocBuffer(payloadSize));
    if (message == nullptr) {
        VmiFreeBuffer(head);
        return nullPair;
    }

    int rc = memcpy_s(message, payloadSize, reinterpret_cast<uint8_t *>(head + 1), payloadSize);
    if (rc != EOK) {
        ERR("memcpy_s error");
        VmiFreeBuffer(head);
        VmiFreeBuffer(message);
        return nullPair;
    }

    VmiFreeBuffer(head);
    return std::pair<uint8_t*, uint32_t>(message, payloadSize);
}

/**
 * @brief: handle first fragment
 * @param [in] packet: packet
 * @return: nullptr, No complete package was received.
 */
std::pair<uint8_t*, uint32_t> MsgReassemble::ProcessFirstFragment(StreamMsgHead& packet)
{
    if (!(m_reassembleQueue.empty())) {
        ERR("REASSEMBLE error, drop some fragment(size %u), this type is %u", m_totalSize, m_type);
        Clear();
    }
    m_totalSize = packet.getPayloadSize();
    m_reassembleQueue.push_back(&packet);
    return std::pair<uint8_t*, uint32_t>(nullptr, 0);
}

/**
 * @brief: handle middle fragment
 * @param [in] packet: packet
 * @return: nullptr, No complete package was received.
 */
std::pair<uint8_t*, uint32_t> MsgReassemble::ProcessMiddleFragment(StreamMsgHead& packet)
{
    std::pair<uint8_t*, uint32_t> nullPair(nullptr, 0);

    if (m_reassembleQueue.empty()) {
        ERR("REASSEMBLE error, drop some fragment(size %u), this type is %u", packet.getPayloadSize(), m_type);
        VmiFreeBuffer(&packet);
    } else {
        if (packet.getPayloadSize() > (MAX_MSG_SIZE - sizeof(StreamMsgHead) - m_totalSize)) {
            ERR("REASSEMBLE error, drop some fragment(size %u), this type is %u", m_totalSize, m_type);
            Clear();
            return nullPair;
        }

        m_totalSize += packet.getPayloadSize();
        m_reassembleQueue.push_back(&packet);
    }
    return nullPair;
}

/**
 * @brief: handle end fragment
 * @param [in] packet: packet
 * @return: Reassemble packet if reassemble ok, or nullptr
 */
std::pair<uint8_t*, uint32_t> MsgReassemble::ProcessEndFragment(StreamMsgHead& packet)
{
    std::pair<uint8_t*, uint32_t> nullPair(nullptr, 0);
    if (m_reassembleQueue.empty()) {
        ERR("REASSEMBLE error, drop some fragment(size %u), this type is %u", packet.getPayloadSize(), m_type);
        VmiFreeBuffer(&packet);
        return nullPair;
    }

    if (packet.getPayloadSize() > (MAX_MSG_SIZE - sizeof(StreamMsgHead) - m_totalSize)) {
        ERR("REASSEMBLE error, drop some fragment(size %u), this type is %u", m_totalSize, m_type);
        Clear();
        return nullPair;
    }

    m_reassembleQueue.push_back(&packet);
    m_totalSize += packet.getPayloadSize();

    uint8_t *message = reinterpret_cast<uint8_t *>(VmiAllocBuffer(m_totalSize));
    if (message == nullptr) {
        ERR("REASSEMBLE alloc buf failed, drop some fragment(size %u), this type is %u", m_totalSize, m_type);
        Clear();
        return nullPair;
    }

    uint8_t* messageData = message;
    while (!(m_reassembleQueue.empty())) {
        StreamMsgHead* tmp = m_reassembleQueue.front();
        int cpyError = memcpy_s(
            messageData, m_totalSize, reinterpret_cast<uint8_t *>(tmp + 1), tmp->getPayloadSize());
        if (cpyError != EOK) {
            ERR("memcpy_s error");
            VmiFreeBuffer(message);
            Clear();
            return nullPair;
        }
        m_reassembleQueue.pop_front();
        messageData += tmp->getPayloadSize();
        VmiFreeBuffer(tmp);
    }

    return std::pair<uint8_t*, uint32_t>(message, m_totalSize);
}

/**
 * @brief: packet reassemble
 * @param [in] packet: packet
 * @return: Reassemble packet if reassemble ok, or nullptr
 */
std::pair<uint8_t*, uint32_t> MsgReassemble::Reassemble(StreamMsgHead& packet)
{
    std::pair<uint8_t*, uint32_t> resultMsg(nullptr, 0);

    if (!CheckPacket(packet)) {
        return resultMsg;
    }

    switch (packet.getFlag()) {
        case SINGLE_FRAGMENT: {
            resultMsg = ProcessSingleFragment(packet);
            break;
        }
        case FIRST_FRAGMENT: {
            resultMsg = ProcessFirstFragment(packet);
            break;
        }
        case MIDDLE_FRAGMENT: {
            resultMsg = ProcessMiddleFragment(packet);
            break;
        }
        case END_FRAGMENT: {
            resultMsg = ProcessEndFragment(packet);
            Clear();
            break;
        }
    }

    return resultMsg;
}
